import os
from hashlib import sha256
from typing import Dict

import jieba
from dotenv import load_dotenv
from langchain_chroma import Chroma
from langchain_community.llms.tongyi import Tongyi
from langchain_community.retrievers import BM25Retriever
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate


def format_history(history, max_epoch=3):
    # 每轮对话有 用户问题 和 助手回复
    if len(history) > 2 * max_epoch:
        history = history[-2 * max_epoch:]
    return "\n".join([f"{i['role']}：{i['content']}" for i in history])


def format_docs(docs: list[Document]) -> str:
    """拼接多个 Document 的 page_content """
    return "\n\n".join(doc.page_content for doc in docs)


def get_retriever(k=20, embedding_model=None):
    """获取向量数据库的检索器"""

    # 1、初始化 Chroma 客户端
    vectorstore = Chroma(
        persist_directory="vectorstore",
        embedding_function=embedding_model,
    )

    # 2、创建向量数据库检索器
    retriever = vectorstore.as_retriever(
        search_type="similarity",  # 检索方式，similarity 或 mmr
        search_kwargs={"k": k},
    )

    return retriever


def get_llm():
    # 大模型
    load_dotenv()
    TONGYI_API_KEY = os.getenv("TONGYI_API_KEY")
    llm = Tongyi(model="qwen-turbo", api_key=TONGYI_API_KEY)
    return llm

# 创造多查询函数，运用语言模型在问题进行重述后执行多查询
def multi_query(input: Dict[str,str],llm,multi_query_num)-> list[str]:

    #1、构建提示词模板
    prompt=PromptTemplate(
        input_variables=["query","multi_query_num"],
        template="""
        请根据提供的原始问题，生成{multi_query_num}个查询问题，每个问题结束后使用换行符与下一个问题分隔。
        请勿返回重复的查询。
        原始问题：{query}
        """
    )

    #2、构造查询链条
    multi_chain=(
        {
            "query": lambda x: x.get("query"),
            "multi_query_num": lambda x: x.get("multi_query_num"),
        }
        | prompt
        | llm
        | StrOutputParser()
        | (lambda x: print(f"===== 多查询结果: {x}=====") or x)
        | (lambda x: [item.strip() for item in x.split("\n") if item.strip()]) #一定要规范返回的问题，不然问题中会出现一些意外的标点符号
    )

    #3、执行链条查询得到的查询结果为列表
    query_list=multi_chain.invoke({"query": input.get("query"),"multi_query_num": multi_query_num})

    return query_list



# 构造函数用于去除重复文档，进行文档重排序
def remove_duplicate_docs(docs: list[list[Document]],k,number) -> list[Document]:
    #1、创建一个紫癜用于存储分数
    docs_sore={}

    #2、创建一个列表用于存储文档
    not_duplicate_docs = {}

    #3、遍历传入的二维文档列表，使用哈希函数进行去重和记录每个文档的分数
    for doc in docs:
     for rank,content in enumerate(doc):
        key=sha256(content.page_content.encode("utf-8")).hexdigest()
        not_duplicate_docs[key]=content
        docs_sore[key]=docs_sore.get(key,0)+1/(k+rank) #检索出来的文档列表是根据评分进行排序的，0，1.。。排序越靠前评分越高

    #4、根据分数字典，对记录的文档重新排序，返回一个排序后的列表
    docs_sore=sorted(docs_sore.items(),key=lambda x:x[1],reverse=True)
    rerank_docs=[]
    for key,value in docs_sore:
        rerank_docs.append(not_duplicate_docs[key])
    return rerank_docs[:number]


def rephrase_retrieve(input: Dict[str, str], llm, retriever,multi_query_num):
    """重述用户query，检索向量数据库"""

    # 1、重述query的prompt
    rephrase_prompt = PromptTemplate.from_template(
        """
        根据对话历史简要完善最新的用户消息，使其更加具体。只输出完善后的问题。如果问题不需要完善，请直接输出原始问题。

        {history}
        用户：{query}
        """
    )

    # 2、重述链条：根据历史和当前 query 生成更具体问题
    rephrase_chain = (
            {
                "history": lambda x: format_history(x.get("history")),
                "query": lambda x: x.get("query"),
            }
            | rephrase_prompt
            | llm
            | StrOutputParser()
            | (lambda x: print(f"===== 重述后的查询: {x}=====") or x)
            | (lambda x: x.split("\n"))

    )

    # 3、执行重述和多查询和关键词查询bm25
    rephrase_query = rephrase_chain.invoke({"history": input.get("history"), "query": input.get("query")})
    query_multi_list=multi_query({"query":rephrase_query},llm,multi_query_num)
    query_keywords_list=create_hy_retriever.invoke(rephrase_query,k=20)

    # 4、使用重述、多查询后的query进行检索
    double_list = []
    for query in query_multi_list:
        double_list.append(retriever.invoke(query, k=3))


    retrieve_result_list=remove_duplicate_docs(double_list,k=1,number=10)
    print(f"===== 检索出来的文档: {retrieve_result_list}=====")
    return retrieve_result_list


def  get_rag_chain(retrieve_result, llm):
    """构建RAG链条：使用检索结果、历史记录、用户查询，提交大模型生成回复"""
    # 1、Prompt 模板
    prompt = PromptTemplate(
        input_variables=["context", "history", "query"],
        template="""
    你是一个专业的中文问答助手，擅长基于提供的资料详细的回答问题。
    请仅根据以下背景资料以及历史消息回答问题，如无法找到答案，请直接回答“我不知道”。

    背景资料：{context}

    历史消息：[{history}]

    问题：{query}

    回答：""",
    )

    # 2、定义 RAG 链条
    rag_chain = (
            {
                "context": lambda x: format_docs(retrieve_result),
                "history": lambda x: format_history(x.get("history")),
                "query": lambda x: x.get("query"),
            }
            | prompt
            | llm
            | StrOutputParser()  # 输出解析器，将输出解析为字符串
    )

    return rag_chain

